import depthai as dai
import cv2
import numpy as np
from datetime import datetime

class VideoRecorder:
    def __init__(self):
        # 初始化视频录制参数
        self.resolution = (1280, 720)
        self.fps = 30
        self.output_file = f"recording_{datetime.now().strftime('%Y%m%d_%H%M%S')}.avi"
        
        # 初始化视频编码器
        self.fourcc = cv2.VideoWriter_fourcc(*'XVID')  # 可尝试其他编码器：MJPG, H264 (需安装对应编码器)
        self.video_writer = cv2.VideoWriter(
            self.output_file,
            self.fourcc,
            self.fps,
            self.resolution
        )
        
        # DepthAI初始化
        try:
            self.device = dai.Device()
            self.pipeline = self.create_pipeline()
            self.device.startPipeline(self.pipeline)
            self.rgb_queue = self.device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
            print(f"🎥 开始录制 {self.resolution[0]}x{self.resolution[1]}@{self.fps}fps → {self.output_file}")
        except Exception as e:
            print(f"设备初始化失败: {e}")
            raise

    def create_pipeline(self):
        pipeline = dai.Pipeline()

        # RGB摄像头配置
        cam_rgb = pipeline.create(dai.node.ColorCamera)
        cam_rgb.setBoardSocket(dai.CameraBoardSocket.CAM_A)
        cam_rgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
        cam_rgb.setFps(self.fps)
        cam_rgb.setIspScale(2, 3)  # 1280x720

        # 输出配置
        xout_rgb = pipeline.createXLinkOut()
        xout_rgb.setStreamName("rgb")
        cam_rgb.isp.link(xout_rgb.input)

        return pipeline

    def record(self):
        try:
            while True:
                # 获取RGB帧
                rgb_data = self.rgb_queue.tryGet()
                if rgb_data:
                    frame = rgb_data.getCvFrame()
                    
                    # 写入视频文件
                    self.video_writer.write(frame)
                    
                    # 实时预览（可选）
                    cv2.imshow("Recording Preview", frame)
                    if cv2.waitKey(1) == ord('q'):
                        break

        except KeyboardInterrupt:
            print("\n录制终止")
        finally:
            # 释放资源
            self.video_writer.release()
            cv2.destroyAllWindows()
            print(f"✅ 视频已保存至: {self.output_file}")

if __name__ == '__main__':
    recorder = VideoRecorder()
    recorder.record()
